use dftd4::prelude::*;

fn main_test() {
    // atom indices
    let numbers = vec![1, 6, 6, 6, 6, 6, 6, 1, 1, 1, 1, 1, 35, 6, 9, 9, 9];
    // geometry in angstrom
    #[rustfmt::skip]
    let positions = vec![
        0.002144194,   0.361043475,   0.029799709,
        0.015020592,   0.274789738,   1.107648016,
        1.227632658,   0.296655040,   1.794629427,
        1.243958826,   0.183702791,   3.183703934,
        0.047958213,   0.048915002,   3.886484583,
       -1.165135654,   0.026954348,   3.200213281,
       -1.181832083,   0.139828643,   1.810376587,
        2.155807907,   0.399177037,   1.249441585,
        2.184979344,   0.198598553,   3.716170761,
        0.060934662,  -0.040672756,   4.964014252,
       -2.093220602,  -0.078628959,   3.745125056,
       -2.122845437,   0.123257119,   1.277645797,
       -0.268325907,  -3.194209024,   1.994458950,
        0.049999933,  -5.089197474,   1.929391171,
        0.078949601,  -5.512441335,   0.671851563,
        1.211983937,  -5.383996300,   2.498664481,
       -0.909987405,  -5.743747328,   2.570721738,
    ];
    // convert angstrom to bohr
    let positions = positions.iter().map(|&x| x / 0.52917721067).collect::<Vec<f64>>();
    // generate DFTD4 model
    let model = DFTD4Model::new(&numbers, &positions, None, None, None);
    // retrive the DFTD4 parameters
    let param = DFTD4Param::load_rational_damping("b97m", true);
    // obtain the dispersion energy and gradient, without sigma
    let (energy, gradient, _) = model.get_dispersion(&param, true).into();
    let gradient = gradient.unwrap();
    println!("Dispersion energy: {}", energy);
    println!("Dispersion gradient:");
    gradient.chunks(3).for_each(|chunk| println!("{:16.9?}", chunk));

    #[rustfmt::skip]
    let gradient_ref = vec![
         2.98598566e-06,  5.58662750e-05, -2.26040542e-04,
         1.75816159e-05,  3.76346114e-04, -5.59737686e-04,
         5.30037419e-04,  3.39528029e-04, -2.58687563e-04,
         5.28657666e-04,  2.71291979e-04,  3.25797992e-04,
         2.86698935e-05,  2.42198161e-04,  6.35784343e-04,
        -4.99037369e-04,  2.74343468e-04,  3.55926359e-04,
        -5.08919994e-04,  3.36480299e-04, -2.53165937e-04,
         1.90422032e-04,  3.91738111e-05, -1.05867702e-04,
         1.88137796e-04,  1.62741316e-05,  1.07526135e-04,
         6.73432246e-06,  2.93353505e-06,  2.23028509e-04,
        -1.97182105e-04,  1.82664904e-05,  1.26683921e-04,
        -2.07136902e-04,  4.89015922e-05, -1.12337721e-04,
        -2.01762905e-04, -1.23133705e-03, -2.26387771e-04,
         7.59912392e-06, -1.01235495e-04, -7.86554908e-06,
         3.02563455e-05, -2.19783887e-04, -2.31677736e-04,
         2.31670341e-04, -2.08064229e-04,  9.53780006e-05,
        -1.48713265e-04, -2.61183219e-04,  1.11642946e-04,
    ];

    let l2_diff =
        gradient.iter().zip(gradient_ref.iter()).map(|(x, y)| (x - y).powi(2)).sum::<f64>().sqrt();
    println!("L2 difference: {:16.8e}", l2_diff);
    assert!(l2_diff < 1e-9);
}

#[test]
fn test() {
    main_test();
}

fn main() {
    main_test();
}

/* equivalent PySCF with dftd4

// https://github.com/dftd4/dftd4/blob/bb9bd9459566a4199ea1a73488008c27ad029ecb/python/dftd4/test_pyscf.py#L87-L133

```python
import pyscf
from pyscf import gto
import dftd4.pyscf as disp

mol = gto.M(atom="""
    H    0.002144194   0.361043475   0.029799709
    C    0.015020592   0.274789738   1.107648016
    C    1.227632658   0.296655040   1.794629427
    C    1.243958826   0.183702791   3.183703934
    C    0.047958213   0.048915002   3.886484583
    C   -1.165135654   0.026954348   3.200213281
    C   -1.181832083   0.139828643   1.810376587
    H    2.155807907   0.399177037   1.249441585
    H    2.184979344   0.198598553   3.716170761
    H    0.060934662  -0.040672756   4.964014252
    H   -2.093220602  -0.078628959   3.745125056
    H   -2.122845437   0.123257119   1.277645797
    Br  -0.268325907  -3.194209024   1.994458950
    C    0.049999933  -5.089197474   1.929391171
    F    0.078949601  -5.512441335   0.671851563
    F    1.211983937  -5.383996300   2.498664481
    F   -0.909987405  -5.743747328   2.570721738
""")

d4 = disp.DFTD4Dispersion(mol, xc="b97m")
print(d4.kernel()[1])
```

*/
